package com.tianji.aigc.memory;

import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.tianji.aigc.entity.ChatDetail;
import com.tianji.aigc.service.ChatDetailService;
import com.tianji.common.utils.CollUtils;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;

import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.List;

/**
 * Mysql实现会话存储
 *
 * @author zgp
 * @version V1.0
 * @date 2025-04-21 20:47
 */
@Slf4j
public class MysqlChatMemory implements ChatMemory {
    @Resource
    private ChatDetailService chatDetailService;

    @Override
    public void add(String conversationId, List<Message> messages) {
        if (CollUtils.isEmpty(messages)) {
            return;
        }
        List<ChatDetail> chatDetails = new ArrayList<>();
        // String[] userIdAndConversationId = conversationId.split("_");
        messages.stream().map(message -> {
            ChatDetail chatDetail = new ChatDetail();
            chatDetail.setConversationId(conversationId);
            chatDetail.setOriginalMessage(MessageUtil.toJson(message));
            chatDetail.setCreateTime(LocalDateTime.now());
            return chatDetail;
        }).forEach(chatDetails::add);
        chatDetailService.saveBatch(chatDetails);
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        if (lastN < 0) {
            return List.of();
        }
        List<ChatDetail> chatDetails = chatDetailService.lambdaQuery().eq(ChatDetail::getConversationId, conversationId).last("limit 0 , " + lastN).list();
        if (CollUtils.isEmpty(chatDetails)) {
            return List.of();
        }
        return chatDetails
                .stream()
                .map(chatDetail -> MessageUtil.toMessage(chatDetail.getOriginalMessage()))
                .toList();
    }

    @Override
    public void clear(String conversationId) {
        chatDetailService.remove(Wrappers.<ChatDetail>lambdaQuery().eq(ChatDetail::getConversationId, conversationId));
    }

    /**
     * 根据对话ID优化对话记录，删除最后的2条消息，因为这2条消息是从路由智能体存储的，请求由后续的智能体处理
     * 为了确保历史消息的完整性，所以需要将中间转发的消息清理掉
     *
     * @param conversationId 对话的唯一标识符
     */
    public void optimization(String conversationId) {
        // var redisKey = this.getKey(conversationId);
        // var listOps = this.stringRedisTemplate.boundListOps(redisKey);
        // // 从Redis列表右侧弹出2个元素
        // listOps.rightPop(2);
        log.info("conversationId = {}", conversationId);
        List<Long> deleteIds = chatDetailService.lambdaQuery()
                .eq(ChatDetail::getConversationId, conversationId)
                // 根据创建时间降序，区最新的两条数据
                .orderByDesc(ChatDetail::getCreateTime)
                .last("limit 0 , 2")
                .list()
                .stream()
                .map(chatDetail -> chatDetail.getId())
                .toList();
        chatDetailService.removeBatchByIds(deleteIds);
    }

}
